{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 在这个代码中，seq2seq无法计算出梯度"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 加载数据集"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Go.\tVa !\n",
      "Hi.\tSalut !\n",
      "Run!\tCours !\n",
      "Run!\tCourez !\n",
      "Who?\tQui ?\n",
      "Wow!\tÇa alors !\n",
      "\n"
     ]
    }
   ],
   "source": [
    "from tqdm.notebook import tqdm\n",
    "import os\n",
    "import mindspore as ms\n",
    "import mindspore.ops as ops\n",
    "import numpy as np\n",
    "import mindspore.nn as nn\n",
    "import random\n",
    "\n",
    "def read_data_nmt(data_path=\"./fra-eng/\"):\n",
    "    \"\"\"载入“英语－法语”数据集\"\"\"\n",
    "    with open(os.path.join(data_path, 'fra.txt'), 'r',\n",
    "             encoding='utf-8') as f:\n",
    "        return f.read()\n",
    "\n",
    "def preprocess_nmt(text:str):\n",
    "    \"\"\"文本预处理\"\"\"\n",
    "    # 替换空白字符为普通空格，并转为全小写\n",
    "    text = text.replace('\\u202f', ' ').replace('\\xa0', ' ').lower()\n",
    "    def no_space(char, prev_char):\n",
    "        return char in set(',.!?') and prev_char != ' '\n",
    "    out = [\" \" + char if k>0 and no_space(char, text[k-1]) else char for k, char in enumerate(tqdm(text))]\n",
    "    return \"\".join(out)\n",
    "\n",
    "raw_text = read_data_nmt()\n",
    "print(raw_text[:75])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f651d399bb1f497bb08b5874aa5e122e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/11489286 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "go .\tva !\n",
      "hi .\tsalut !\n",
      "run !\tcours !\n",
      "run !\tcourez !\n",
      "who ?\tqui ?\n",
      "wow !\tça alors !\n"
     ]
    }
   ],
   "source": [
    "text = preprocess_nmt(raw_text)\n",
    "print(text[:80])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "([['go', '.'],\n",
       "  ['hi', '.'],\n",
       "  ['run', '!'],\n",
       "  ['run', '!'],\n",
       "  ['who', '?'],\n",
       "  ['wow', '!']],\n",
       " '...',\n",
       " [['va', '!'],\n",
       "  ['salut', '!'],\n",
       "  ['cours', '!'],\n",
       "  ['courez', '!'],\n",
       "  ['qui', '?'],\n",
       "  ['ça', 'alors', '!']])"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def tokenize_nmt(text:str, num_examples=None):\n",
    "    \"\"\"词元化“英语－法语”数据数据集\"\"\"\n",
    "    source, target = [], []\n",
    "    for k, line in enumerate(text.split('\\n')):\n",
    "        if (num_examples is not None) and (k > num_examples):\n",
    "            break\n",
    "        pair = line.split('\\t')\n",
    "        if len(pair) == 2:\n",
    "            source.append(pair[0].split(\" \"))\n",
    "            target.append(pair[1].split(\" \"))\n",
    "    return source, target\n",
    "\n",
    "s, t = tokenize_nmt(text)\n",
    "s[:6], \"...\", t[:6]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "import collections\n",
    "class Vocab:\n",
    "    \"\"\"一个词汇表的实现\"\"\"\n",
    "    def __init__(self, tokens:list, min_freq=0, reserved_tokens:list=None) -> None:\n",
    "        self.default_index = None\n",
    "        if tokens is not None:\n",
    "            # 当第一个条件满足时，就不会跳到第二个判断，避免了空列表报错的情况。\n",
    "            if len(tokens)!=0 and isinstance(tokens[0], list):\n",
    "                tokens = [i for line in tokens for i in line]\n",
    "        else:\n",
    "            tokens = []\n",
    "        if reserved_tokens is None:\n",
    "            reserved_tokens = []\n",
    "        counter=collections.Counter(tokens)\n",
    "        # 按出现词频从高到低排序\n",
    "        self._token_freqs = sorted(counter.items(), key=lambda x:x[1], reverse=True)\n",
    "        # 通过列表,利用序号访问词元。\n",
    "        self.idx_to_token = [] + reserved_tokens # 未知词元<unk>的索引为0, 保留词元排在最前\n",
    "        self.token_to_idx = {\n",
    "            i: k\n",
    "            for k, i in enumerate(self.idx_to_token) \n",
    "        }\n",
    "        \n",
    "        for token, freq in self._token_freqs:\n",
    "            if freq < min_freq:  # 过滤掉出现频率低于要求的词\n",
    "                break\n",
    "            if token not in self.token_to_idx:  \n",
    "                self.idx_to_token.append(token)\n",
    "                self.token_to_idx[token] = len(self.idx_to_token) - 1\n",
    "        \n",
    "    def __len__(self):\n",
    "        return len(self.idx_to_token)\n",
    "    \n",
    "    def __getitem__(self, input_tokens):\n",
    "        \"\"\"输入单字串或序列, 将其全部转化为序号编码\"\"\"\n",
    "        if isinstance(input_tokens, str):\n",
    "            out =  self.token_to_idx.get(input_tokens, self.default_index)\n",
    "            if out is None:\n",
    "                raise Exception('Please call \"set_default_index\" before getting unknown index')\n",
    "            return out\n",
    "        return [self.__getitem__(token) for token in input_tokens]\n",
    "    \n",
    "    def __repr__(self) -> str:\n",
    "        show_items = 5 if len(self) > 5 else len(self)\n",
    "        out = f\"<Vocab with {len(self)} tokens: \"\n",
    "        for i in range(show_items):\n",
    "            out += f'\"{self.idx_to_token[i]}\", '\n",
    "        out += \"...>\"\n",
    "        return out\n",
    "\n",
    "    def __contains__(self, token:str) -> bool:\n",
    "        return token in self.idx_to_token\n",
    "\n",
    "    def to_tokens(self, input_keys):\n",
    "        \"\"\"输入单s索引或序列, 将其全部转化为词元\"\"\"\n",
    "        if isinstance(input_keys, int):\n",
    "            return self.idx_to_token[input_keys] if input_keys < len(self) else self.idx_to_token[0]\n",
    "        elif isinstance(input_keys, (list, tuple)):\n",
    "            return [self.to_tokens(keys) for keys in input_keys]\n",
    "        else:\n",
    "            return self.idx_to_token[0]\n",
    "    \n",
    "    def get_itos(self):\n",
    "        return self.idx_to_token\n",
    "    \n",
    "    def get_stoi(self):\n",
    "        return self.token_to_idx\n",
    "    \n",
    "    def set_default_index(self, idx):\n",
    "        if isinstance(idx, int):\n",
    "            self.default_index = idx\n",
    "        else:\n",
    "            raise Exception(f\"Only type int allowed, got {type(idx)}\")\n",
    "    \n",
    "    def lookup_indices(self, input_tokens):\n",
    "        return self.__getitem__(input_tokens)\n",
    "    \n",
    "    def lookup_tokens(self, idx):\n",
    "        return self.to_tokens(idx)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [],
   "source": [
    "random.shuffle(t)\n",
    "random.shuffle(s)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 初始化Vocab"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {},
   "outputs": [],
   "source": [
    "min_freq = 2\n",
    "unk_token = \"<unk>\"\n",
    "pad_token = \"<pad>\"\n",
    "sos_token = \"<sos>\"\n",
    "eos_token = \"<eos>\"\n",
    "\n",
    "special_tokens = [\n",
    "    unk_token,\n",
    "    pad_token,\n",
    "    sos_token,\n",
    "    eos_token,\n",
    "]\n",
    "en_vocab = Vocab(s, min_freq=2, \n",
    "                        reserved_tokens=special_tokens)\n",
    "de_vocab = Vocab(t, min_freq=2, \n",
    "                        reserved_tokens=special_tokens)\n",
    "\n",
    "len(en_vocab), len(de_vocab)\n",
    "assert en_vocab[unk_token] == de_vocab[unk_token]\n",
    "assert en_vocab[pad_token] == de_vocab[pad_token]\n",
    "unk_index = en_vocab[unk_token]\n",
    "pad_index = en_vocab[pad_token]\n",
    "en_vocab.set_default_index(unk_index)\n",
    "de_vocab.set_default_index(unk_index)\n",
    "\n",
    "# 取出测试用数据\n",
    "test_en = en_vocab.lookup_indices(s[0:128])\n",
    "test_de = de_vocab.lookup_indices(t[0:128])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [],
   "source": [
    "def pad_sequence(sequences:list, padding_value:int):\n",
    "    '''将序列填充到等长并返回mindspore张量'''\n",
    "    # Find the length of the longest sequence in the batch\n",
    "    max_length = max(len(seq) for seq in sequences)\n",
    "    padded_sequences = ops.full((len(sequences), max_length), padding_value, dtype=ms.int64)\n",
    "    # Copy the sequences into the padded array\n",
    "    for i, seq in enumerate(sequences):\n",
    "        padded_sequences[i, :len(seq)] = ms.tensor(seq).astype(np.int64)\n",
    "    # 换轴，保证输出为时序优先\n",
    "    padded_sequences = padded_sequences.swapaxes(0, 1)\n",
    "    return padded_sequences  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "((19, 128), (19, 128))"
      ]
     },
     "execution_count": 51,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data_en = pad_sequence(test_en, pad_index)\n",
    "data_de = pad_sequence(test_en, pad_index)\n",
    "\n",
    "data_en.shape, data_de.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Tensor(shape=[19], dtype=Int64, value= [  30,  146,   19, 1615,    4,    1,    1,    1,    1,    1,    1,    1,    1,    1,    1,    1,    1,    1,    1])"
      ]
     },
     "execution_count": 52,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data_en[:, 5]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 99,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Encoder(nn.Cell):\n",
    "    def __init__(self, input_dim, embedding_dim, hidden_dim, n_layers, dropout):\n",
    "        super().__init__()\n",
    "        self.hidden_dim = hidden_dim\n",
    "        self.n_layers = n_layers\n",
    "        self.embedding = nn.Embedding(input_dim, embedding_dim)\n",
    "        self.rnn = nn.LSTM(embedding_dim, hidden_dim, n_layers, dropout=dropout)\n",
    "        self.dropout = nn.Dropout(p=dropout)\n",
    "        \n",
    "    def construct(self, src):\n",
    "        # src = [src length, batch size]\n",
    "        embedded = self.dropout(self.embedding(src))\n",
    "        # embedded = [src length, batch size, embedding dim]\n",
    "        outputs, (hidden, cell) = self.rnn(embedded)\n",
    "        # outputs = [src length, batch size, hidden dim * n directions]\n",
    "        # hidden = [n layers * n directions, batch size, hidden dim]\n",
    "        # cell = [n layers * n directions, batch size, hidden dim]\n",
    "        # outputs are always from the top hidden layer\n",
    "        return hidden, cell\n",
    "\n",
    "class Decoder(nn.Cell):\n",
    "    def __init__(self, output_dim, embedding_dim, hidden_dim, n_layers, dropout):\n",
    "        super().__init__()\n",
    "        self.output_dim = output_dim\n",
    "        self.hidden_dim = hidden_dim\n",
    "        self.n_layers = n_layers\n",
    "        self.embedding = nn.Embedding(output_dim, embedding_dim)\n",
    "        self.rnn = nn.LSTM(embedding_dim, hidden_dim, n_layers, dropout=dropout)\n",
    "        self.fc_out = nn.Dense(hidden_dim, output_dim)\n",
    "        self.dropout = nn.Dropout(p=dropout)\n",
    "        \n",
    "    def construct(self, input:ms.Tensor, hidden:ms.Tensor, cell:ms.Tensor):\n",
    "        # input = [batch size]\n",
    "        # hidden = [n layers * n directions, batch size, hidden dim]\n",
    "        # cell = [n layers * n directions, batch size, hidden dim]\n",
    "        # n directions in the decoder will both always be 1, therefore:\n",
    "        # hidden = [n layers, batch size, hidden dim]\n",
    "        # context = [n layers, batch size, hidden dim]\n",
    "        input = input.unsqueeze(0)\n",
    "        # input = [1, batch size]\n",
    "        embedded = self.dropout(self.embedding(input))\n",
    "        # embedded = [1, batch size, embedding dim]\n",
    "        output, (hidden, cell) = self.rnn(embedded, (hidden, cell))\n",
    "        # output = [seq length, batch size, hidden dim * n directions]\n",
    "        # hidden = [n layers * n directions, batch size, hidden dim]\n",
    "        # cell = [n layers * n directions, batch size, hidden dim]\n",
    "        # seq length and n directions will always be 1 in this decoder, therefore:\n",
    "        # output = [1, batch size, hidden dim]\n",
    "        # hidden = [n layers, batch size, hidden dim]\n",
    "        # cell = [n layers, batch size, hidden dim]\n",
    "        prediction = self.fc_out(output.squeeze(0))\n",
    "        # prediction = [batch size, output dim]\n",
    "        return prediction, hidden, cell"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 建立Seq2Seq模型。方法一导致梯度丢失"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 100,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "class Seq2Seq(nn.Cell):\n",
    "    def __init__(self, encoder:Encoder, decoder:Decoder):\n",
    "        super().__init__()\n",
    "        self.encoder = encoder\n",
    "        self.decoder = decoder\n",
    "        assert encoder.hidden_dim == decoder.hidden_dim, \\\n",
    "            \"Hidden dimensions of encoder and decoder must be equal!\"\n",
    "        assert encoder.n_layers == decoder.n_layers, \\\n",
    "            \"Encoder and decoder must have equal number of layers!\"\n",
    "        \n",
    "    def construct(self, src:ms.Tensor, trg:ms.Tensor, teacher_forcing_ratio):\n",
    "        # src = [src length, batch size]\n",
    "        # trg = [trg length, batch size]\n",
    "        # teacher_forcing_ratio is probability to use teacher forcing\n",
    "        # e.g. if teacher_forcing_ratio is 0.75 we use ground-truth inputs 75% of the time\n",
    "        batch_size = trg.shape[1]\n",
    "        trg_length = trg.shape[0]\n",
    "        trg_vocab_size = self.decoder.output_dim\n",
    "        # tensor to store decoder outputs\n",
    "        # outputs = ops.zeros((trg_length, batch_size, trg_vocab_size))\n",
    "        # outputs = []\n",
    "        outputs = None\n",
    "        # last hidden state of the encoder is used as the initial hidden state of the decoder\n",
    "        hidden, cell = self.encoder(src)\n",
    "        # hidden = [n layers * n directions, batch size, hidden dim]\n",
    "        # cell = [n layers * n directions, batch size, hidden dim]\n",
    "        # first input to the decoder is the <sos> tokens\n",
    "        input = trg[0,:]\n",
    "        # input = [batch size]\n",
    "        for t in range(1, trg_length):\n",
    "            # insert input token embedding, previous hidden and previous cell states\n",
    "            # receive output tensor (predictions) and new hidden and cell states\n",
    "            output, hidden, cell = self.decoder(input, hidden, cell)\n",
    "            if outputs is None:\n",
    "                outputs = ops.zeros((trg_length, batch_size, trg_vocab_size), dtype=output.dtype)\n",
    "            # output = [batch size, output dim]\n",
    "            # hidden = [n layers, batch size, hidden dim]\n",
    "            # cell = [n layers, batch size, hidden dim]\n",
    "            # place predictions in a tensor holding predictions for each token\n",
    "            # outputs.append(output)\n",
    "            outputs[t] = output\n",
    "            # decide if we are going to use teacher forcing or not\n",
    "            teacher_force = random.random() < teacher_forcing_ratio\n",
    "            # get the highest predicted token from our predictions\n",
    "            top1 = output.argmax(1)\n",
    "            # if teacher forcing, use actual next token as next input\n",
    "            # if not, use predicted token\n",
    "            input = trg[t] if teacher_force else top1\n",
    "            # input = [batch size]\n",
    "        return outputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 101,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_dim = len(de_vocab)\n",
    "output_dim = len(en_vocab)\n",
    "encoder_embedding_dim = 256\n",
    "decoder_embedding_dim = 256\n",
    "hidden_dim = 512\n",
    "n_layers = 2\n",
    "encoder_dropout = 0.5\n",
    "decoder_dropout = 0.5\n",
    "\n",
    "encoder = Encoder(\n",
    "    input_dim,\n",
    "    encoder_embedding_dim,\n",
    "    hidden_dim,\n",
    "    n_layers,\n",
    "    encoder_dropout,\n",
    ")\n",
    "\n",
    "decoder = Decoder(\n",
    "    output_dim,\n",
    "    decoder_embedding_dim,\n",
    "    hidden_dim,\n",
    "    n_layers,\n",
    "    decoder_dropout,\n",
    ")\n",
    "\n",
    "model = Seq2Seq(encoder, decoder)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 102,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(19, 128) (19, 128)\n",
      "[ 220  312   96   57 2662  154    4   30  519   16    6    4    1    1\n",
      "    1    1    1    1    1] \n",
      " [ 220  312   96   57 2662  154    4   30  519   16    6    4    1    1\n",
      "    1    1    1    1    1]\n",
      "['quelques', 'vrai', 'on', 'votre', 'tenez-vous', \"j'étais\", '.', 'pour', 'désolée', 'tom', 'de', '.', '<pad>', '<pad>', '<pad>', '<pad>', '<pad>', '<pad>', '<pad>']\n",
      "[\"you've\", 'both', 'been', 'very', 'impressive', 'today', '.', \"i'm\", 'proud', 'of', 'you', '.', '<pad>', '<pad>', '<pad>', '<pad>', '<pad>', '<pad>', '<pad>']\n",
      "(19, 128, 10012)\n",
      "(2304, 10012) (2304,)\n",
      "9.210166 (Tensor(shape=[17851, 256], dtype=Float32, value=\n",
      "[[ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " ...\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00]]), Tensor(shape=[2048, 256], dtype=Float32, value=\n",
      "[[ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " ...\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00]]), Tensor(shape=[2048, 512], dtype=Float32, value=\n",
      "[[ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " ...\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00]]), Tensor(shape=[2048, 512], dtype=Float32, value=\n",
      "[[ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " ...\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00]]), Tensor(shape=[2048, 512], dtype=Float32, value=\n",
      "[[ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " ...\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00]]), Tensor(shape=[2048], dtype=Float32, value= [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00]), Tensor(shape=[2048], dtype=Float32, value= [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00]), Tensor(shape=[2048], dtype=Float32, value= [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00]), Tensor(shape=[2048], dtype=Float32, value= [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00]), Tensor(shape=[10012, 256], dtype=Float32, value=\n",
      "[[ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " ...\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00]]), Tensor(shape=[2048, 256], dtype=Float32, value=\n",
      "[[ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " ...\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00]]), Tensor(shape=[2048, 512], dtype=Float32, value=\n",
      "[[ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " ...\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00]]), Tensor(shape=[2048, 512], dtype=Float32, value=\n",
      "[[ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " ...\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00]]), Tensor(shape=[2048, 512], dtype=Float32, value=\n",
      "[[ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " ...\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00]]), Tensor(shape=[2048], dtype=Float32, value= [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00]), Tensor(shape=[2048], dtype=Float32, value= [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00]), Tensor(shape=[2048], dtype=Float32, value= [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00]), Tensor(shape=[2048], dtype=Float32, value= [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00]), Tensor(shape=[10012, 512], dtype=Float32, value=\n",
      "[[ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " ...\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00]]), Tensor(shape=[10012], dtype=Float32, value= [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00]))\n"
     ]
    }
   ],
   "source": [
    "criterion = nn.CrossEntropyLoss(ignore_index=pad_index)\n",
    "\n",
    "def forward_fn(src, trg):\n",
    "    # src = [src length, batch size]\n",
    "    # trg = [trg length, batch size]\n",
    "    output = model(src, trg, 0.5)\n",
    "    print(output.shape)\n",
    "    # output = [trg length, batch size, trg vocab size]\n",
    "    output_dim = output.shape[-1]\n",
    "    output = output[1:].view(-1, output_dim)\n",
    "    # output = [(trg length - 1) * batch size, trg vocab size]\n",
    "    trg = trg[1:].view(-1)\n",
    "    # trg = [(trg length - 1) * batch size]\n",
    "    print(output.shape, trg.shape)\n",
    "    loss = criterion(output, trg.astype(ms.int32))\n",
    "    return loss\n",
    "\n",
    "grad_fn = ms.value_and_grad(forward_fn, grad_position=None, weights=model.trainable_params())\n",
    "\n",
    "\n",
    "src = data_en\n",
    "trg = data_de\n",
    "# 检查输入数据\n",
    "print(src.shape, trg.shape)  # 数据形状\n",
    "print(src[:, 0], \"\\n\", trg[:, 0])  # 第一个句子具体内容\n",
    "print(de_vocab.lookup_tokens(src[:, 0].tolist()))  # 文本化\n",
    "print(en_vocab.lookup_tokens(trg[:, 0].tolist()))\n",
    "loss, grad = grad_fn(src, trg)  # 求导\n",
    "print(loss, grad)  # 显示梯度"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 尝试2 创建全零数组，然后将结果加上去，仍然梯度数据丢失"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 91,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Seq2Seq(nn.Cell):\n",
    "    def __init__(self, encoder:Encoder, decoder:Decoder):\n",
    "        super().__init__()\n",
    "        self.encoder = encoder\n",
    "        self.decoder = decoder\n",
    "        assert encoder.hidden_dim == decoder.hidden_dim, \\\n",
    "            \"Hidden dimensions of encoder and decoder must be equal!\"\n",
    "        assert encoder.n_layers == decoder.n_layers, \\\n",
    "            \"Encoder and decoder must have equal number of layers!\"\n",
    "        \n",
    "    def construct(self, src:ms.Tensor, trg:ms.Tensor, teacher_forcing_ratio):\n",
    "        # src = [src length, batch size]\n",
    "        # trg = [trg length, batch size]\n",
    "        # teacher_forcing_ratio is probability to use teacher forcing\n",
    "        # e.g. if teacher_forcing_ratio is 0.75 we use ground-truth inputs 75% of the time\n",
    "        batch_size = trg.shape[1]\n",
    "        trg_length = trg.shape[0]\n",
    "        trg_vocab_size = self.decoder.output_dim\n",
    "        # tensor to store decoder outputs\n",
    "        outputs = ops.zeros((trg_length, batch_size, trg_vocab_size))\n",
    "        # outputs = []\n",
    "\n",
    "        # last hidden state of the encoder is used as the initial hidden state of the decoder\n",
    "        hidden, cell = self.encoder(src)\n",
    "        # hidden = [n layers * n directions, batch size, hidden dim]\n",
    "        # cell = [n layers * n directions, batch size, hidden dim]\n",
    "        # first input to the decoder is the <sos> tokens\n",
    "        input = trg[0,:]\n",
    "        # input = [batch size]\n",
    "        for t in range(1, trg_length):\n",
    "            # insert input token embedding, previous hidden and previous cell states\n",
    "            # receive output tensor (predictions) and new hidden and cell states\n",
    "            output, hidden, cell = self.decoder(input, hidden, cell)\n",
    "            # output = [batch size, output dim]\n",
    "            # hidden = [n layers, batch size, hidden dim]\n",
    "            # cell = [n layers, batch size, hidden dim]\n",
    "            # place predictions in a tensor holding predictions for each token\n",
    "            # outputs.append(output)\n",
    "            outputs[t] += output\n",
    "            # decide if we are going to use teacher forcing or not\n",
    "            teacher_force = random.random() < teacher_forcing_ratio\n",
    "            # get the highest predicted token from our predictions\n",
    "            top1 = output.argmax(1)\n",
    "            # if teacher forcing, use actual next token as next input\n",
    "            # if not, use predicted token\n",
    "            input = trg[t] if teacher_force else top1\n",
    "            # input = [batch size]\n",
    "        return outputs\n",
    "\n",
    "input_dim = len(de_vocab)\n",
    "output_dim = len(en_vocab)\n",
    "encoder_embedding_dim = 256\n",
    "decoder_embedding_dim = 256\n",
    "hidden_dim = 512\n",
    "n_layers = 2\n",
    "encoder_dropout = 0.5\n",
    "decoder_dropout = 0.5\n",
    "\n",
    "encoder = Encoder(\n",
    "    input_dim,\n",
    "    encoder_embedding_dim,\n",
    "    hidden_dim,\n",
    "    n_layers,\n",
    "    encoder_dropout,\n",
    ")\n",
    "\n",
    "decoder = Decoder(\n",
    "    output_dim,\n",
    "    decoder_embedding_dim,\n",
    "    hidden_dim,\n",
    "    n_layers,\n",
    "    decoder_dropout,\n",
    ")\n",
    "\n",
    "model = Seq2Seq(encoder, decoder)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 92,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(19, 128) (19, 128)\n",
      "[ 220  312   96   57 2662  154    4   30  519   16    6    4    1    1\n",
      "    1    1    1    1    1] \n",
      " [ 220  312   96   57 2662  154    4   30  519   16    6    4    1    1\n",
      "    1    1    1    1    1]\n",
      "['quelques', 'vrai', 'on', 'votre', 'tenez-vous', \"j'étais\", '.', 'pour', 'désolée', 'tom', 'de', '.', '<pad>', '<pad>', '<pad>', '<pad>', '<pad>', '<pad>', '<pad>']\n",
      "[\"you've\", 'both', 'been', 'very', 'impressive', 'today', '.', \"i'm\", 'proud', 'of', 'you', '.', '<pad>', '<pad>', '<pad>', '<pad>', '<pad>', '<pad>', '<pad>']\n",
      "(19, 128, 10012)\n",
      "(2304, 10012) (2304,)\n",
      "9.202961 (Tensor(shape=[17851, 256], dtype=Float32, value=\n",
      "[[ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " ...\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00]]), Tensor(shape=[2048, 256], dtype=Float32, value=\n",
      "[[ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " ...\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00]]), Tensor(shape=[2048, 512], dtype=Float32, value=\n",
      "[[ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " ...\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00]]), Tensor(shape=[2048, 512], dtype=Float32, value=\n",
      "[[ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " ...\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00]]), Tensor(shape=[2048, 512], dtype=Float32, value=\n",
      "[[ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " ...\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00]]), Tensor(shape=[2048], dtype=Float32, value= [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00]), Tensor(shape=[2048], dtype=Float32, value= [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00]), Tensor(shape=[2048], dtype=Float32, value= [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00]), Tensor(shape=[2048], dtype=Float32, value= [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00]), Tensor(shape=[10012, 256], dtype=Float32, value=\n",
      "[[ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " ...\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00]]), Tensor(shape=[2048, 256], dtype=Float32, value=\n",
      "[[ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " ...\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00]]), Tensor(shape=[2048, 512], dtype=Float32, value=\n",
      "[[ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " ...\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00]]), Tensor(shape=[2048, 512], dtype=Float32, value=\n",
      "[[ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " ...\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00]]), Tensor(shape=[2048, 512], dtype=Float32, value=\n",
      "[[ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " ...\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00]]), Tensor(shape=[2048], dtype=Float32, value= [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00]), Tensor(shape=[2048], dtype=Float32, value= [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00]), Tensor(shape=[2048], dtype=Float32, value= [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00]), Tensor(shape=[2048], dtype=Float32, value= [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00]), Tensor(shape=[10012, 512], dtype=Float32, value=\n",
      "[[ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " ...\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00]]), Tensor(shape=[10012], dtype=Float32, value= [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00]))\n"
     ]
    }
   ],
   "source": [
    "criterion = nn.CrossEntropyLoss(ignore_index=pad_index)\n",
    "\n",
    "def forward_fn(src, trg):\n",
    "    # src = [src length, batch size]\n",
    "    # trg = [trg length, batch size]\n",
    "    output = model(src, trg, 0.5)\n",
    "    print(output.shape)\n",
    "    # output = [trg length, batch size, trg vocab size]\n",
    "    output_dim = output.shape[-1]\n",
    "    output = output[1:].view(-1, output_dim)\n",
    "    # output = [(trg length - 1) * batch size, trg vocab size]\n",
    "    trg = trg[1:].view(-1)\n",
    "    # trg = [(trg length - 1) * batch size]\n",
    "    print(output.shape, trg.shape)\n",
    "    loss = criterion(output, trg.astype(ms.int32))\n",
    "    return loss\n",
    "\n",
    "grad_fn = ms.value_and_grad(forward_fn, grad_position=None, weights=model.trainable_params())\n",
    "\n",
    "\n",
    "src = data_en\n",
    "trg = data_de\n",
    "# 检查输入数据\n",
    "print(src.shape, trg.shape)  # 数据形状\n",
    "print(src[:, 0], \"\\n\", trg[:, 0])  # 第一个句子具体内容\n",
    "print(de_vocab.lookup_tokens(src[:, 0].tolist()))  # 文本化\n",
    "print(en_vocab.lookup_tokens(trg[:, 0].tolist()))\n",
    "loss, grad = grad_fn(src, trg)  # 求导\n",
    "print(loss, grad)  # 显示梯度"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 尝试3 将输出存入list内，最后将list变为张量，梯度出现"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 97,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Seq2Seq(nn.Cell):\n",
    "    def __init__(self, encoder:Encoder, decoder:Decoder):\n",
    "        super().__init__()\n",
    "        self.encoder = encoder\n",
    "        self.decoder = decoder\n",
    "        assert encoder.hidden_dim == decoder.hidden_dim, \\\n",
    "            \"Hidden dimensions of encoder and decoder must be equal!\"\n",
    "        assert encoder.n_layers == decoder.n_layers, \\\n",
    "            \"Encoder and decoder must have equal number of layers!\"\n",
    "        \n",
    "    def construct(self, src:ms.Tensor, trg:ms.Tensor, teacher_forcing_ratio):\n",
    "        # src = [src length, batch size]\n",
    "        # trg = [trg length, batch size]\n",
    "        # teacher_forcing_ratio is probability to use teacher forcing\n",
    "        # e.g. if teacher_forcing_ratio is 0.75 we use ground-truth inputs 75% of the time\n",
    "        batch_size = trg.shape[1]\n",
    "        trg_length = trg.shape[0]\n",
    "        trg_vocab_size = self.decoder.output_dim\n",
    "        # tensor to store decoder outputs\n",
    "        # outputs = ops.zeros((trg_length, batch_size, trg_vocab_size))\n",
    "        outputs = []\n",
    "\n",
    "        # last hidden state of the encoder is used as the initial hidden state of the decoder\n",
    "        hidden, cell = self.encoder(src)\n",
    "        # hidden = [n layers * n directions, batch size, hidden dim]\n",
    "        # cell = [n layers * n directions, batch size, hidden dim]\n",
    "        # first input to the decoder is the <sos> tokens\n",
    "        input = trg[0,:]\n",
    "        # input = [batch size]\n",
    "        for t in range(1, trg_length):\n",
    "            # insert input token embedding, previous hidden and previous cell states\n",
    "            # receive output tensor (predictions) and new hidden and cell states\n",
    "            output, hidden, cell = self.decoder(input, hidden, cell)\n",
    "            if len(outputs) == 0:\n",
    "                outputs.append(ops.zeros(output.shape, dtype=output.dtype))\n",
    "            # output = [batch size, output dim]\n",
    "            # hidden = [n layers, batch size, hidden dim]\n",
    "            # cell = [n layers, batch size, hidden dim]\n",
    "            # place predictions in a tensor holding predictions for each token\n",
    "            outputs.append(output)\n",
    "            # outputs[t] += output\n",
    "            # decide if we are going to use teacher forcing or not\n",
    "            teacher_force = random.random() < teacher_forcing_ratio\n",
    "            # get the highest predicted token from our predictions\n",
    "            top1 = output.argmax(1)\n",
    "            # if teacher forcing, use actual next token as next input\n",
    "            # if not, use predicted token\n",
    "            input = trg[t] if teacher_force else top1\n",
    "            # input = [batch size]\n",
    "        return ops.stack(outputs, axis=0)\n",
    "\n",
    "input_dim = len(de_vocab)\n",
    "output_dim = len(en_vocab)\n",
    "encoder_embedding_dim = 256\n",
    "decoder_embedding_dim = 256\n",
    "hidden_dim = 512\n",
    "n_layers = 2\n",
    "encoder_dropout = 0.5\n",
    "decoder_dropout = 0.5\n",
    "\n",
    "encoder = Encoder(\n",
    "    input_dim,\n",
    "    encoder_embedding_dim,\n",
    "    hidden_dim,\n",
    "    n_layers,\n",
    "    encoder_dropout,\n",
    ")\n",
    "\n",
    "decoder = Decoder(\n",
    "    output_dim,\n",
    "    decoder_embedding_dim,\n",
    "    hidden_dim,\n",
    "    n_layers,\n",
    "    decoder_dropout,\n",
    ")\n",
    "\n",
    "model = Seq2Seq(encoder, decoder)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 98,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(19, 128) (19, 128)\n",
      "[ 220  312   96   57 2662  154    4   30  519   16    6    4    1    1\n",
      "    1    1    1    1    1] \n",
      " [ 220  312   96   57 2662  154    4   30  519   16    6    4    1    1\n",
      "    1    1    1    1    1]\n",
      "['quelques', 'vrai', 'on', 'votre', 'tenez-vous', \"j'étais\", '.', 'pour', 'désolée', 'tom', 'de', '.', '<pad>', '<pad>', '<pad>', '<pad>', '<pad>', '<pad>', '<pad>']\n",
      "[\"you've\", 'both', 'been', 'very', 'impressive', 'today', '.', \"i'm\", 'proud', 'of', 'you', '.', '<pad>', '<pad>', '<pad>', '<pad>', '<pad>', '<pad>', '<pad>']\n",
      "(19, 128, 10012)\n",
      "(2304, 10012) (2304,)\n",
      "9.212521 (Tensor(shape=[17851, 256], dtype=Float32, value=\n",
      "[[-1.46045203e-08,  3.51113911e-08, -4.73580073e-08 ... -1.56414082e-08,  1.63767555e-09,  2.24970123e-08],\n",
      " [ 1.23067302e-05, -8.12171402e-06, -4.73449036e-05 ... -3.91903013e-05, -3.32813543e-05,  1.81453597e-05],\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " ...\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00]]), Tensor(shape=[2048, 256], dtype=Float32, value=\n",
      "[[-2.52703997e-10,  1.61023538e-11, -6.90008467e-11 ... -3.10729116e-11,  3.31097955e-10, -3.07207676e-10],\n",
      " [ 5.32399946e-09, -2.16897900e-09, -8.86855922e-10 ...  3.13205706e-09, -7.24379490e-09,  6.47701004e-09],\n",
      " [ 1.48889141e-08, -7.14810255e-09, -3.97053634e-09 ...  8.88256402e-09, -2.21772360e-08,  1.60532370e-08],\n",
      " ...\n",
      " [-3.86739663e-09,  1.81981763e-09,  9.36064892e-10 ... -2.32282305e-09,  5.73111203e-09, -4.34361258e-09],\n",
      " [ 1.78504607e-08, -8.53754845e-09, -4.38737358e-09 ...  1.06052349e-08, -2.63743765e-08,  1.99078780e-08],\n",
      " [ 1.07420330e-08, -4.98496178e-09, -2.09889328e-09 ...  6.16713125e-09, -1.53736490e-08,  1.22262840e-08]]), Tensor(shape=[2048, 512], dtype=Float32, value=\n",
      "[[-6.07062702e-08, -1.01758936e-07, -5.63267797e-08 ...  6.56859243e-08, -6.27162393e-08, -1.50908789e-07],\n",
      " [ 7.13503994e-08,  1.21963623e-07,  6.55431478e-08 ... -7.65924497e-08,  7.66800312e-08,  1.78899683e-07],\n",
      " [-2.44783482e-09, -2.64000466e-09, -3.18341176e-09 ...  3.56608343e-09, -1.23440347e-09, -5.12763121e-09],\n",
      " ...\n",
      " [-3.02021297e-09, -5.23410515e-09, -2.86763413e-09 ...  3.29765482e-09, -3.36760531e-09, -7.73735032e-09],\n",
      " [-5.87122306e-09, -9.99609284e-09, -5.40163470e-09 ...  6.37004094e-09, -6.27591401e-09, -1.47603822e-08],\n",
      " [-2.70593592e-09, -4.58146276e-09, -2.41383602e-09 ...  2.88262814e-09, -2.95269942e-09, -6.89051438e-09]]), Tensor(shape=[2048, 512], dtype=Float32, value=\n",
      "[[-9.87569193e-11, -2.39650910e-10, -1.26460384e-10 ...  9.51843881e-11, -1.17145210e-10, -2.13739290e-10],\n",
      " [ 4.52382842e-09,  7.53246532e-09,  4.34865211e-09 ... -5.09401410e-09,  4.51507365e-09,  1.11558025e-08],\n",
      " [ 1.35166571e-08,  2.29797177e-08,  1.23001183e-08 ... -1.43976679e-08,  1.42566634e-08,  3.38680586e-08],\n",
      " ...\n",
      " [-3.52235485e-09, -5.88759930e-09, -3.25466920e-09 ...  3.84393983e-09, -3.64465969e-09, -8.72873951e-09],\n",
      " [ 1.59655595e-08,  2.68732041e-08,  1.46963854e-08 ... -1.73204242e-08,  1.66871885e-08,  3.97756139e-08],\n",
      " [ 9.29790822e-09,  1.57564539e-08,  8.64575522e-09 ... -1.00455235e-08,  9.66990665e-09,  2.30826931e-08]]), Tensor(shape=[2048, 512], dtype=Float32, value=\n",
      "[[ 2.95988968e-07, -3.85373426e-07, -2.34464849e-07 ... -1.03798612e-07, -5.52839730e-08,  1.23129311e-07],\n",
      " [-3.52688119e-07,  4.58936000e-07,  2.78897062e-07 ...  1.23506865e-07,  6.53857981e-08, -1.46556118e-07],\n",
      " [ 8.47946513e-09, -1.07754259e-08, -7.67731034e-09 ... -3.21955884e-09, -1.64516600e-09,  3.54926399e-09],\n",
      " ...\n",
      " [ 1.52377950e-08, -1.97936227e-08, -1.20812134e-08 ... -5.33094147e-09, -2.81864665e-09,  6.32870956e-09],\n",
      " [ 2.89741013e-08, -3.77380260e-08, -2.29477184e-08 ... -1.01821893e-08, -5.38148104e-09,  1.20770451e-08],\n",
      " [ 1.36542004e-08, -1.77854407e-08, -1.07251656e-08 ... -4.77230655e-09, -2.54112376e-09,  5.69998404e-09]]), Tensor(shape=[2048], dtype=Float32, value= [-1.11038503e-08,  5.11570420e-07,  1.57242187e-06 ... -3.99442683e-07,  1.83785585e-06,  1.06392497e-06]), Tensor(shape=[2048], dtype=Float32, value= [-6.96384859e-06,  8.30035788e-06, -1.99581137e-07 ... -3.58351542e-07, -6.81440156e-07, -3.20088134e-07]), Tensor(shape=[2048], dtype=Float32, value= [-1.11038503e-08,  5.11570420e-07,  1.57242187e-06 ... -3.99442683e-07,  1.83785585e-06,  1.06392497e-06]), Tensor(shape=[2048], dtype=Float32, value= [-6.96384859e-06,  8.30035788e-06, -1.99581137e-07 ... -3.58351542e-07, -6.81440156e-07, -3.20088134e-07]), Tensor(shape=[10012, 256], dtype=Float32, value=\n",
      "[[ 8.73838474e-07, -4.79070593e-07,  6.91240757e-07 ... -9.03735099e-07,  9.03316163e-07, -1.63388245e-06],\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " ...\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00],\n",
      " [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00 ...  0.00000000e+00,  0.00000000e+00,  0.00000000e+00]]), Tensor(shape=[2048, 256], dtype=Float32, value=\n",
      "[[ 4.81836260e-08,  4.00838438e-08,  8.60568861e-09 ... -7.59630936e-08,  6.91286886e-08, -4.20315764e-08],\n",
      " [-1.81972801e-08, -1.09268292e-08, -1.21626726e-08 ...  2.85439636e-08, -2.41728930e-08,  9.91423210e-09],\n",
      " [-4.90932734e-08, -6.17887821e-08, -5.25002726e-08 ...  6.52712870e-08, -4.42173018e-08,  3.24694120e-08],\n",
      " ...\n",
      " [-2.31232544e-09, -9.98924676e-09,  1.78870625e-08 ... -2.20800267e-10, -1.17070886e-08,  8.21069435e-09],\n",
      " [ 6.04056041e-08,  6.17032043e-08,  1.09154206e-08 ... -9.37301650e-08,  8.50284394e-08, -5.48244863e-08],\n",
      " [ 2.20971934e-08, -7.68980435e-09,  1.61913825e-08 ... -2.96602671e-08,  3.57477070e-08, -6.23550189e-09]]), Tensor(shape=[2048, 512], dtype=Float32, value=\n",
      "[[-5.66594224e-07,  4.41169220e-07, -8.06138473e-07 ... -6.44855504e-07,  9.66424977e-07,  1.01130797e-06],\n",
      " [-1.08593909e-07,  5.34940376e-08, -1.41223026e-07 ... -9.27394055e-08,  1.44972702e-07,  1.56842134e-07],\n",
      " [ 1.22640005e-08, -1.51303841e-08,  2.22868444e-08 ...  1.96610355e-08, -2.94487350e-08, -2.80989614e-08],\n",
      " ...\n",
      " [-5.16111470e-07,  3.35582882e-07, -6.71952307e-07 ... -5.09318056e-07,  7.92099570e-07,  8.59803038e-07],\n",
      " [ 2.94096107e-08,  3.67215591e-10,  6.71825404e-08 ...  1.78741644e-08, -2.68586788e-08, -6.57817623e-09],\n",
      " [-3.24749792e-08,  6.13595432e-08, -6.06392589e-08 ... -7.12337851e-08,  1.02842449e-07,  9.71110836e-08]]), Tensor(shape=[2048, 512], dtype=Float32, value=\n",
      "[[-6.24074517e-08,  7.03610397e-08, -1.10160016e-07 ... -9.57459036e-08,  1.40933153e-07,  1.43018269e-07],\n",
      " [ 1.29692541e-08, -3.93617938e-08,  3.50612233e-08 ...  4.55048763e-08, -6.41116813e-08, -6.77103031e-08],\n",
      " [ 2.94156823e-08, -1.61883889e-07,  1.14087456e-07 ...  1.75739700e-07, -2.47462566e-07, -2.60688040e-07],\n",
      " ...\n",
      " [ 2.27975896e-08, -9.25065002e-09,  3.30454561e-08 ...  1.88759408e-08, -3.06710817e-08, -2.26045103e-08],\n",
      " [-8.53841726e-08,  1.18581063e-07, -1.65596148e-07 ... -1.56432293e-07,  2.29860774e-07,  2.23207095e-07],\n",
      " [-2.84591328e-09,  2.28825758e-08, -1.94355572e-08 ... -2.62329678e-08,  3.78723186e-08,  3.38179120e-08]]), Tensor(shape=[2048, 512], dtype=Float32, value=\n",
      "[[ 1.45657054e-06,  1.43916793e-07,  1.53871156e-07 ... -1.02198328e-06,  5.51993423e-07, -4.74120185e-07],\n",
      " [ 3.67528742e-07, -9.65226477e-08, -4.66437271e-08 ... -2.27494638e-07,  8.81264981e-08, -5.09402263e-08],\n",
      " [-1.87257179e-08, -2.16076899e-08, -1.42705510e-08 ...  1.97128927e-08, -1.68810423e-08,  1.81410815e-08],\n",
      " ...\n",
      " [ 1.49688947e-06, -9.03753801e-08,  4.64510297e-09 ... -9.68815925e-07,  4.50008883e-07, -3.44021856e-07],\n",
      " [-2.03002941e-07,  1.57634190e-07,  9.25461379e-08 ...  1.25509459e-07, -3.33171712e-08, -6.28190122e-09],\n",
      " [-1.86827762e-08,  1.52234108e-07,  9.49281329e-08 ... -2.51735273e-08,  5.47953647e-08, -7.64483943e-08]]), Tensor(shape=[2048], dtype=Float32, value= [ 8.79773233e-06, -4.06426125e-06, -1.56117640e-05 ... -1.64883295e-06,  1.39863887e-05,  2.18948117e-06]), Tensor(shape=[2048], dtype=Float32, value= [ 6.13492302e-05,  9.33747378e-06, -1.78052619e-06 ...  5.09529127e-05, -1.29589716e-06,  6.43278145e-06]), Tensor(shape=[2048], dtype=Float32, value= [ 8.79773233e-06, -4.06426125e-06, -1.56117640e-05 ... -1.64883295e-06,  1.39863887e-05,  2.18948117e-06]), Tensor(shape=[2048], dtype=Float32, value= [ 6.13492302e-05,  9.33747378e-06, -1.78052619e-06 ...  5.09529127e-05, -1.29589716e-06,  6.43278145e-06]), Tensor(shape=[10012, 512], dtype=Float32, value=\n",
      "[[-1.48340085e-04,  1.16452629e-05,  6.56949567e-07 ...  9.40284226e-05, -4.26434381e-05,  3.17244776e-05],\n",
      " [ 2.16483045e-06,  3.21774849e-07,  2.97724313e-07 ... -1.61196681e-06,  9.34333002e-07, -8.13471786e-07],\n",
      " [ 2.15459750e-06,  3.18686773e-07,  2.95366931e-07 ... -1.60349236e-06,  9.28768145e-07, -8.08346556e-07],\n",
      " ...\n",
      " [ 2.32236721e-06,  3.43575863e-07,  3.18394626e-07 ... -1.72843249e-06,  1.00119496e-06, -8.71402221e-07],\n",
      " [ 2.11494171e-06,  3.09232064e-07,  2.87603456e-07 ... -1.57239151e-06,  9.09511130e-07, -7.91013974e-07],\n",
      " [ 2.19892672e-06,  3.24491339e-07,  3.00944407e-07 ... -1.63619200e-06,  9.47474064e-07, -8.24509016e-07]]), Tensor(shape=[10012], dtype=Float32, value= [-4.87178331e-03,  9.88722022e-05,  9.83108330e-05 ...  1.05972467e-04,  9.63050479e-05,  1.00294354e-04]))\n"
     ]
    }
   ],
   "source": [
    "criterion = nn.CrossEntropyLoss(ignore_index=pad_index)\n",
    "\n",
    "def forward_fn(src, trg):\n",
    "    # src = [src length, batch size]\n",
    "    # trg = [trg length, batch size]\n",
    "    output = model(src, trg, 0.5)\n",
    "    print(output.shape)\n",
    "    # output = [trg length, batch size, trg vocab size]\n",
    "    output_dim = output.shape[-1]\n",
    "    output = output[1:].view(-1, output_dim)\n",
    "    # output = [(trg length - 1) * batch size, trg vocab size]\n",
    "    trg = trg[1:].view(-1)\n",
    "    # trg = [(trg length - 1) * batch size]\n",
    "    print(output.shape, trg.shape)\n",
    "    loss = criterion(output, trg.astype(ms.int32))\n",
    "    return loss\n",
    "\n",
    "grad_fn = ms.value_and_grad(forward_fn, grad_position=None, weights=model.trainable_params())\n",
    "\n",
    "\n",
    "src = data_en\n",
    "trg = data_de\n",
    "# 检查输入数据\n",
    "print(src.shape, trg.shape)  # 数据形状\n",
    "print(src[:, 0], \"\\n\", trg[:, 0])  # 第一个句子具体内容\n",
    "print(de_vocab.lookup_tokens(src[:, 0].tolist()))  # 文本化\n",
    "print(en_vocab.lookup_tokens(trg[:, 0].tolist()))\n",
    "loss, grad = grad_fn(src, trg)  # 求导\n",
    "print(loss, grad)  # 显示梯度"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
